#!/usr/bin/env python
#################################################################################
# Copyright 2018 ROBOTIS CO., LTD.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################

# Authors: Gilbert #

import rospy
import os
import json
import numpy as np
import random
import time
import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
from collections import deque
from std_msgs.msg import Float32MultiArray
from src.turtlebot3_dqn.simulation_stage_1_environment import Env
from keras.models import Sequential, load_model
from keras.optimizers import RMSprop
from keras.layers import Dense, Dropout, Activation, SimpleRNN, LSTM
from keras.utils import plot_model
from gazebo_msgs.msg import ModelStates, ModelState
from geometry_msgs.msg import Pose


EPISODES = 3000

class ReinforceAgent():
    def __init__(self, state_size, action_size):
        self.pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        self.dirPath = os.path.dirname(os.path.realpath(__file__))
        self.dirPath = self.dirPath.replace('turtlebot3_dqn/nodes', 'turtlebot3_dqn/save_model/stage_1_')
        self.result = Float32MultiArray()

        self.load_model = False
        self.load_episode = 0
        self.state_size = state_size
        self.action_size = action_size
        self.episode_step = 6000
        self.target_update = 2000
        self.discount_factor = 0.99
        self.learning_rate = 0.00025
        self.epsilon = 1.0
        self.epsilon_decay = 0.99
        self.epsilon_min = 0.05
        self.batch_size = 64
        self.train_start = 64
        self.memory = deque(maxlen=1000000)
        self.start_time = time.time()

        self.model = self.buildModel()
        self.target_model = self.buildModel()

        self.updateTargetModel()

        self.turtlebot = rospy.Publisher('gazebo/set_model_state', ModelState, queue_size=1)

        if self.load_model:
            self.model.set_weights(load_model(self.dirPath+str(self.load_episode)+".h5").get_weights())

            with open(self.dirPath+str(self.load_episode)+'.json') as outfile:
                param = json.load(outfile)
                self.epsilon = param.get('epsilon')

    def buildModel(self):
        # lstm sequential model
        model = Sequential()
        # model.add(Dense(64, input_shape=(self.state_size,), activation='relu', kernel_initializer='lecun_uniform'))
        model.add(LSTM(64, input_shape=(10,364), return_sequences=True))
        model.add(LSTM(64))
        # model.add(Dense(64, input_shape=(self.state_size,), activation='relu', kernel_initializer='lecun_uniform'))
        # model.add(SimpleRNN(1, return_sequences=False, stateful = False))
        # model.add(Dense(64, activation='relu', kernel_initializer='lecun_uniform'))

        model.add(Dense(self.action_size, kernel_initializer='lecun_uniform'))
        model.add(Activation('linear'))
        model.compile(loss='mse', optimizer=RMSprop(lr=self.learning_rate, rho=0.9, epsilon=1e-06))
        model.summary()

        plot_model(model, to_file='model.png')
        return model

    def getQvalue(self, reward, next_target, done):
        if done:
            return reward
        else:
            return reward + self.discount_factor * np.amax(next_target)

    def updateTargetModel(self):
        self.target_model.set_weights(self.model.get_weights())

    def getAction(self, state):
        state = np.asarray(state)
        if np.random.rand() <= self.epsilon:
            self.q_value = np.zeros(self.action_size)
            return random.randrange(self.action_size)
        else:
            q_value = self.model.predict(state.reshape((1, 10, 364)))
            # q_value = self.model.predict(state.reshape((1, len(state))))
            self.q_value = q_value
            return np.argmax(q_value[0])

    def appendMemory(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def trainModel(self, target=False):
        # appended experiences
        mini_batch = random.sample(self.memory, self.batch_size)
        z = 0
        while z < 5:
            X_batch = np.empty((0,10, self.state_size), dtype=np.float64)
            Y_batch = np.empty((0, self.action_size), dtype=np.float64)
            Z_batch = np.empty((0, 1), dtype=np.float64)
            t = True
            states = [0,0,0,0,0,0,0,0,0,0]
            next_states = [0,0,0,0,0,0,0,0,0,0]

            while t:
                x = random.randint(1, len(self.memory)-1)
                if x>12:
                    t=False
            i=0
            actions = self.memory[x][1]
            rewards = self.memory[x][2]
            dones = self.memory[x][4]
            while i <10:
                states[i] = self.memory[x-i][0]
                next_states[i] = self.memory[x-i][3]
                i+=1

            states = np.asarray(states)
            next_states = np.asarray(next_states)

            q_value = self.model.predict(states.reshape((1,10, 364)))
            # q_value = self.model.predict(states.reshape((1, len(states))))
            self.q_value = q_value
    
            #dqn
            if target:
                next_target = self.target_model.predict(next_states.reshape((1,10, 364)))
                # next_target = self.target_model.predict(next_states.reshape((1, len(next_states))))

            else:
                next_target = self.model.predict(next_states.reshape((1,10, 364)))
                # next_target = self.model.predict(next_states.reshape((1, len(next_states))))

            next_q_value = self.getQvalue(rewards, next_target, dones)

            X_batch = np.append(X_batch, np.array([states.copy()]), axis=0)
            Y_sample = q_value.copy()

            Y_sample[0][actions] = next_q_value
            Y_batch = np.append(Y_batch, np.array([Y_sample[0]]), axis=0)

            # sizes must fit to LSTM dimensions
            if dones:
                X_batch = np.append(X_batch, np.array([next_states.copy()]), axis=0)
                Y_batch = np.append(Y_batch, np.array([[rewards] * self.action_size]), axis=0)
            # print len(Y_batch)
            # print X_batch.shape
            X_batch = X_batch.reshape((len(Y_batch), 10, 364))

            self.model.fit(X_batch, Y_batch, batch_size=self.batch_size, epochs=1, verbose=0)
            z =+1

if __name__ == '__main__':
    rospy.init_node('turtlebot3_dqn_stage_1')
    pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
    pub_get_action = rospy.Publisher('get_action', Float32MultiArray, queue_size=5)
    result = Float32MultiArray()
    get_action = Float32MultiArray()

    state_size = 364
    action_size = 5

    env = Env(action_size)

    agent = ReinforceAgent(state_size, action_size)
    scores, episodes = [], []
    global_step = 0
    states = [0,0,0,0,0,0,0,0,0,0]


    with open('/home/wangqiang/catkin_ws/src/turtlebot3_machine_learning/turtlebot3_dqn/results_experiments/dqn_simulation_1_64_64_gui_false_time_environment_blablabla.txt', 'w+') as saveFile:
        # saveFile.write(str(times) + '\n')  write each result to new line
        saveFile.write("time_agent_get_action"+","+"time_env_step"+","+"append_memory"+","+"time_train_model" +","+"time_one_episode_step"+","+"time_publish"+","+"data_time"+","+"getstate_time" +","+"time_setreward" + ","+"time_one_step" + '\n')
        saveFile.flush()

        for e in range(agent.load_episode + 1, EPISODES):
            done = False

            time_env_reset_before = time.time()
            state = env.reset()
            time_env_reset_after = time.time()
            # saveFile.write("time_env_reset:"+ str(time_env_reset_after - time_env_reset_before) + '\n')

            score = 0
            # action appending, first random, then start training
            for t in range(agent.episode_step):
                start_time = time.time()
                if global_step>=10:
                    states[0]=states[1]
                    states[1]=states[2]
                    states[2]=states[3]
                    states[3]=states[4]
                    states[4]=states[5]
                    states[5]=states[6]
                    states[6]=states[7]
                    states[7]=states[8]
                    states[8]=states[9]
                    states[9]=state
                    action = agent.getAction(states)
                else:
                    action = 2
                    states[global_step] = state

                next_state, reward, done = env.step(action)
                agent.appendMemory(state, action, reward, next_state, done)

                if len(agent.memory) >= agent.train_start:
                    if global_step <= agent.target_update:
                        agent.trainModel()
                    else:
                        agent.trainModel(True)

                score += reward
                state = next_state
                get_action.data = [action, score, reward]
                pub_get_action.publish(get_action)

                if e % 100 == 0:
                    agent.model.save(agent.dirPath + str(e) + '.h5')
                    with open(agent.dirPath + str(e) + '.json', 'w') as outfile:
                        json.dump(param_dictionary, outfile)

                if t >= 500:
                    rospy.loginfo("Time out!!")
                    done = True

                if done:
                    result.data = [score, np.max(agent.q_value)]
                    pub_result.publish(result)
                    agent.updateTargetModel()
                    scores.append(score)
                    episodes.append(e)
                    m, s = divmod(int(time.time() - start_time), 60)
                    h, m = divmod(m, 60)

                    rospy.loginfo('Ep: %d score: %.2f memory: %d epsilon: %.2f time: %d:%02d:%02d',
                                  e, score, len(agent.memory), agent.epsilon, h, m, s)
                    param_keys = ['epsilon']
                    param_values = [agent.epsilon]
                    param_dictionary = dict(zip(param_keys, param_values))
                    break

                global_step += 1
                if global_step % agent.target_update == 0:
                    rospy.loginfo("UPDATE TARGET NETWORK")

                time_episode_end = time.time()
                # saveFile.write(str(time_get_action_after - time_get_action_before)+","+str(time_env_step_after - time_env_step_before) +","+str(time_append_memory_after - time_append_memory_before)+","+str(time_agent_train_model_after - time_agent_train_model_before)+","+ str(time_episode_end - start_time)+","+str(time_publish)+","+str(data_time)+","+str(getstate_time) +","+str(time_setreward)+","+str(time_step) + '\n')

            if agent.epsilon > agent.epsilon_min:
                agent.epsilon *= agent.epsilon_decay
